import json
from pathlib import Path
import numpy as np
import pandas as pd
from scipy.interpolate import PchipInterpolator


def create_manhattan_masks(map_row_count: int):
    map_col_count = 2 * map_row_count if map_row_count % 2 == 1 else 2 * map_row_count - 1
    # 第一个map_row_count+1对应于距离为0至map_col_count，距离为0的元素没有用
    masks = np.zeros((map_row_count + 1, map_row_count + 1, map_col_count), dtype=bool)
    center_c = (map_col_count - 1) // 2 if map_col_count % 2 == 1 else map_col_count // 2 - 1  # 转换为0-based索引
    # ====生成网格坐标====
    r_indices, c_indices = np.indices((map_row_count + 1, map_col_count))
    # ====计算曼哈顿距离====
    # 小屋在第1行(0-based)
    manhattan_dist = np.abs(r_indices - 0) + np.abs(c_indices - center_c)
    # ====创建mask====
    for dist in range(1, map_row_count + 1):
        masks[dist] = manhattan_dist <= dist
    # 排除第1行(小屋所在行)
    masks[:, 0, :] = False
    mask_areas = np.sum(masks.astype(int), axis=(1, 2))
    return masks, mask_areas


def create_euclidean_masks(map_row_count: int):
    map_col_count = 2 * map_row_count if map_row_count % 2 == 1 else 2 * map_row_count - 1
    # 最高维每个元素是指与农屋最大距离为i的mask，距离为0的元素没有用
    masks = np.zeros((map_row_count + 1, map_row_count + 1, map_col_count), dtype=bool)
    center_c = (map_col_count - 1) // 2 if map_col_count % 2 == 1 else map_col_count // 2 - 1  # 转换为0-based索引
    # ====生成网格坐标====
    r_indices, c_indices = np.indices((map_row_count + 1, map_col_count))
    # ====计算欧几里德距离====
    # 小屋在第1行(0-based)
    euclidean_dist = np.sqrt((r_indices - 0) ** 2 + (c_indices - center_c) ** 2)
    # ====创建mask====
    for dist in range(1, map_row_count + 1):
        masks[dist] = euclidean_dist <= dist
    # 排除第1行(小屋所在行)
    masks[:, 0, :] = False
    mask_areas = np.sum(masks.astype(int), axis=(1, 2))
    return masks, mask_areas


def create_squre_masks(map_row_count: int):
    map_col_count = 2 * map_row_count if map_row_count % 2 == 1 else 2 * map_row_count - 1
    # 最高维每个元素是指与农屋最大距离为i的mask，距离为0的元素没有用
    masks = np.zeros((map_row_count + 1, map_row_count + 1, map_col_count), dtype=bool)
    center_c = (map_col_count - 1) // 2 if map_col_count % 2 == 1 else map_col_count // 2 - 1  # 转换为0-based索引
    for dist in range(1, map_row_count + 1):
        left = center_c - dist if dist % 2 == 1 else center_c - dist + 1
        right = center_c + dist
        masks[dist, 1 : 1 + dist, left : right + 1] = True
    mask_areas = np.sum(masks.astype(int), axis=(1, 2))
    return masks, mask_areas


def calc_areas_of_time_matrix(t_matrix):
    """
    计算农田面积
    :return:最高维是最大距离（1至map_row_count），低维是[曼哈顿面积,曼哈顿时长,欧几里德面积,欧几里德时长,矩形面积,矩形时长]
    """
    map_row_count = t_matrix.shape[0] - 1
    result = np.zeros((map_row_count + 1, 6), dtype=int)
    for index, (masks, mask_areas) in enumerate(
        [
            create_manhattan_masks(map_row_count),
            create_euclidean_masks(map_row_count),
            create_squre_masks(map_row_count),
        ]
    ):
        for dist in range(1, map_row_count + 1):
            result[dist, index * 2], result[dist, index * 2 + 1] = (
                mask_areas[dist],
                np.sum(t_matrix[masks[dist]]),
            )
    return result[1:]


def 工作时长平滑化(df: pd.DataFrame, min_work_time: int, max_work_time: int):
    new_df = pd.DataFrame({"时长": np.arange(min_work_time, max_work_time + 1)})
    for t in ("曼哈顿", "欧几里德", "矩形"):
        # 保持数据单调性的三次样条插值
        pchip_model = PchipInterpolator(df[f"{t}时长"], df[f"{t}面积"])
        new_df[f"{t}面积"] = pchip_model(np.arange(min_work_time, max_work_time + 1))
    new_df["平均面积"] = new_df[["曼哈顿面积", "欧几里德面积", "矩形面积"]].mean(axis=1)
    return new_df


def main():
    T = 16
    DAY_HOUR = 24
    GAME_HOUR_IN_REAL_SEC = 18.75
    map_row_count = 30
    map_col_count = 2 * map_row_count if map_row_count % 2 == 1 else 2 * map_row_count - 1
    center_c = (map_col_count - 1) // 2 if map_col_count % 2 == 1 else map_col_count // 2 - 1  # 转换为0-based索引
    # ========计算走路用时========
    walk_speed = 1.5 * GAME_HOUR_IN_REAL_SEC  # 格/每小时
    # ====欧几里德距离====
    r_indices, c_indices = np.indices((map_row_count + 1, map_col_count))
    # 小屋在第1行(0-based)
    matrix_euclidean_dist = np.sqrt((r_indices - 0) ** 2 + (c_indices - center_c) ** 2)
    # 乘2是因为一来一回
    matrix_time_walk = matrix_euclidean_dist / walk_speed * 2
    # 每天能种T*1.1棵橡树，一共只需要往返一趟，所以走路用时平摊到每棵树上都要除以T*1.1
    matrix_time_walk_tree = matrix_time_walk / (T * 1.1)
    # ================正式开始================
    Path("output").mkdir(exist_ok=True)
    result = {}
    columns = ["曼哈顿面积", "曼哈顿时长", "欧几里德面积", "欧几里德时长", "矩形面积", "矩形时长"]
    tree_data = pd.read_csv("tree_data.csv", header=0, encoding="utf-8-sig", index_col=0)
    # ========计算种树用时========
    # 只需要考虑橡树
    橡树生长周期 = tree_data.loc["橡树", "生长周期d"]
    matrix_time_plant_tree = (
        np.ones((map_row_count + 1, map_col_count), dtype=float) / tree_data.loc["橡树", "种植速度/h"]
    )
    matrix_time_tree = matrix_time_walk_tree + matrix_time_plant_tree
    df_tree = 工作时长平滑化(
        pd.DataFrame(calc_areas_of_time_matrix(matrix_time_tree), columns=columns), 橡树生长周期 * T, 橡树生长周期 * T
    )
    result.update({"种橡树": df_tree.to_dict("records")[0]})
    # ========计算砍树用时========
    matrix_time_chop_tree = (
        np.ones((map_row_count + 1, map_col_count), dtype=float) * tree_data.loc["橡树", "砍伐时长h"]
    )
    # 乘4是因为砍树工人需要分4次搬完8个原木
    matrix_time_chop = matrix_time_walk_tree * (tree_data.loc["橡树", "砍伐获取"] / 2) + matrix_time_chop_tree
    df_chop = 工作时长平滑化(
        pd.DataFrame(calc_areas_of_time_matrix(matrix_time_chop), columns=columns),
        橡树生长周期 * T,
        橡树生长周期 * T * 6,
    )
    df_chop_save_cond = (df_chop["时长"] % 30 == 0) & (df_chop["平均面积"] >= df_tree.loc[0, "平均面积"])
    result.update({"砍橡树": df_chop[df_chop_save_cond].nsmallest(1, "平均面积").to_dict("records")[0]})
    # ========计算种田用时========
    最快作物生长周期d, 最慢作物生长周期d = 4, 10
    农屋工人数量 = 3
    作物播种时长h, 作物收获时长h = 4, 22
    matrix_time_plant_crops = np.ones((map_row_count + 1, map_col_count), dtype=float) * (
        (作物播种时长h + 作物收获时长h) / GAME_HOUR_IN_REAL_SEC
    )
    matrix_time_crops = matrix_time_walk + matrix_time_plant_crops
    df_crops = 工作时长平滑化(
        pd.DataFrame(calc_areas_of_time_matrix(matrix_time_crops), columns=columns),
        最快作物生长周期d * T * 农屋工人数量,
        最慢作物生长周期d * T * 农屋工人数量,
    )
    df_crops[["时长", "平均面积"]].to_csv(
        "output/result_crops.csv", header=True, index=False, encoding="utf-8-sig", float_format="{:.2f}".format
    )
    # ========计算采集用时========
    for tree, fruit in (("枫树", "枫糖浆"), ("栗子树", "栗子"), ("松树", "松脂")):
        收获周期 = tree_data.loc[tree, "收获周期d"]
        matrix_time_gather_fruit = (
            np.ones((map_row_count + 1, map_col_count), dtype=float) * tree_data.loc[tree, "收获时长h"]
        )
        matrix_time_fruit = matrix_time_walk + matrix_time_gather_fruit
        df_fruit = 工作时长平滑化(
            pd.DataFrame(calc_areas_of_time_matrix(matrix_time_fruit), columns=columns), 收获周期 * T, 收获周期 * T
        )
        result.update({f"采集{fruit}": df_fruit.to_dict("records")[0]})
    # ================输出================
    result_df = pd.DataFrame.from_dict(result, orient="index")
    result_df.index.name = "操作类型"
    result_df[["时长", "平均面积"]].to_csv(
        "output/result.csv", header=True, index=True, encoding="utf-8-sig", float_format="{:.2f}".format
    )


main()
